package gold.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Created by fanzhenyu02 on 2020/6/27.
 * Union Find 算法实现
 */
public class UnionFindLarge {
    // 存储一棵树
    private Map<Integer, Integer> parent = new HashMap<>();
    // 记录树的“重量”
    private Map<Integer, Integer> size = new HashMap<>();

    public UnionFindLarge() {

    }

    /* 将 p 和 q 连接 */
    public void union(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        if (rootP == rootQ)
            return;

        // 小树接到大树下面，较平衡
        if (size.getOrDefault(rootP, 0) > size.getOrDefault(rootQ, 0)) {
            parent.put(rootQ, rootP);
            size.put(rootP, size.getOrDefault(rootP, 0) + size.getOrDefault(rootQ, 0));
        } else {
            parent.put(rootP, rootQ);
            size.put(rootQ, size.getOrDefault(rootQ, 0) + size.getOrDefault(rootP, 0));
        }
    }

    /* 判断 p 和 q 是否连通 */
    public boolean connected(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        return rootP == rootQ;
    }

    private int find(int x) {
        while (parent.getOrDefault(x, x).intValue() != x) {
            // 进行路径压缩
            parent.put(x, parent.getOrDefault(x, x));
            x = parent.getOrDefault(x, x);
        }
        return x;
    }

    /* 返回图中有多少个连通分量 */
    public int count() {
        return parent.size();
    }

    /* 返回图中联通集合 */
    public Map<Integer, List<Integer>> outputAggregateSet() {
        Map<Integer, List<Integer>> aggregateMap = new HashMap<>();

        // 遍历所有key
        for (Integer i : parent.keySet()) {
            int parentId = find(i);
            if (!aggregateMap.containsKey(parentId)) aggregateMap.put(parentId, new ArrayList<>());
            aggregateMap.get(parentId).add(i);
        }

        return aggregateMap;
    }

    public static void main(String[] args) {
        /**
         * Given input of (1,2), (2,3), (4,5), (7,8), (5,6,7)
         * Return (1,2,3), (4,5,6,7,8)
         */

        UnionFindLarge uf = new UnionFindLarge();
        uf.union(1, 2);
        uf.union(2, 3);
        uf.union(4, 5);
        uf.union(7, 8);
        uf.union(5, 6);
        uf.union(6, 7);

        final Map<Integer, List<Integer>> integerListMap = uf.outputAggregateSet();
        System.out.println(integerListMap.toString());// 注意输出集合的时候把key也带上
    }

}
